import pymysql
from config import MYSQL_CONFIG  # 修改为相对导入
import logging
from datetime import datetime

class MySQLStorage:
    """MySQL存储类，用于将爬取的新闻保存到MySQL数据库"""
    
    def __init__(self, table_name):
        self.table_name = table_name
        self.conn = None
        self.cursor = None
        
    def connect(self):
        """连接到MySQL数据库"""
        try:
            self.conn = pymysql.connect(**MYSQL_CONFIG)
            self.cursor = self.conn.cursor()
            self.ensure_table_exists()
            return True
        except Exception as e:
            logging.error(f"MySQL连接失败: {str(e)}")
            return False
            
    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
            
    def ensure_table_exists(self):
        """确保数据表存在"""
        try:
            create_table_sql = f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                id BIGINT AUTO_INCREMENT PRIMARY KEY,
                title VARCHAR(255),
                content TEXT,
                url VARCHAR(512),
                publish_time DATETIME,
                crawl_time DATETIME DEFAULT CURRENT_TIMESTAMP
            )
            """
            self.cursor.execute(create_table_sql)
            self.conn.commit()
            logging.info(f"确保数据表 {self.table_name} 存在")
        except Exception as e:
            logging.error(f"创建表失败: {str(e)}")
            raise
        
    def save(self, data):
        """保存数据到MySQL"""
        if not self.connect():
            return False
            
        try:
            # 准备SQL语句
            sql = f"""
            INSERT INTO {self.table_name} (title, content, url, publish_time, crawl_time)
            VALUES (%s, %s, %s, %s, %s)
            """
            
            # 准备数据
            title = data.get('title', '')
            content = data.get('content', '')
            url = data.get('url', '')
            publish_time = data.get('publish_time', datetime.now())
            crawl_time = datetime.now()
            
            # 执行SQL
            self.cursor.execute(sql, (title, content, url, publish_time, crawl_time))
            self.conn.commit()
            
            logging.info(f"成功保存到MySQL: {title}")
            return True
            
        except Exception as e:
            logging.error(f"保存到MySQL失败: {str(e)}")
            self.conn.rollback()
            return False
            
        finally:
            self.close()
        
    # ... 其余代码保持不变 ... 